import collections
import cPickle
import getpass
import os
import subprocess
import sys
import tempfile
import time

from batchcat.exceptions import JobSubmitFail
from batchcat import jobscripts

class JobCountFailed(Exception):
    pass

class ResourceManager(object):
    def __init__(self, name, default_queue, unthrottled_limit, 
                 statepath='bcat_resman_state.pkl'):
        self.name = name
        self.queue = default_queue
        self.unthrottled_limit = unthrottled_limit

        self._joblist = []

        self.pause = 5
                
        self.jobs_running = 0

    def submit(self):
        self._joblist.reverse()
        njobs_total = len(self._joblist) 

        while len(self._joblist) > 0:
            try:
                jobs_running = self.running_job_count()
            except JobCountFailed:
                msg = "Failed to get job count; setting pause to {} seconds"
                warnings.warn(msg.format(self.pause*2))
                self.pause *= 2
            else:
                self.pause = 5

            njobs_to_submit = self.unthrottled_limit - jobs_running
            if njobs_to_submit > len(self._joblist):
                njobs_to_submit = len(self._joblist)

            for i in xrange(njobs_to_submit):
                jobscr = self._joblist.pop()
                with tempfile.NamedTemporaryFile() as fp:
                    fp.write(str(jobscr))
                    fp.flush()
                                        
                    self.state[jobscr] = self._submit(fp.name)
                    

            self.save_state()
            time.sleep(self.pause)                
            sys.stdout.write(
                "{0}/{1} jobs submitted\r".format(njobs_total - 
                                                  len(self._joblist),
                                                  njobs_total))
            sys.stdout.flush()


        print "{0}/{0} jobs submitted\r".format(njobs_total - 
                                                len(self._joblist),
                                                njobs_total)
        return njobs_total

    def _submit(self, path):
        return 'submitted'

    def running_job_count(self):
        raise NotImplementedError("Abstract base class not for end use")

    def add_job(self):
        pass
    
    def add_script(self, scrobj):
        self._joblist.append(scrobj)
    
    def save_state(self):
        with open(self.statepath, 'w') as fp:
            cPickle.dump(self.state)

    def write(self, path=None):
        if path is None:
            path = os.getcwd()
        elif os.access(path, os.F_OK | os.W_OK):
            if not os.path.isdir(path):
                msg = "Cannot write scripts: path `{}' exists " + \
                    "and is not a directory"
                raise IOError(msg.format(path))
        else:
            os.mkdir(path)

        tot_count = 0
        
        nscripts = len(self._joblist)
        for script_obj in self._joblist:
            if tot_count % 100 == 0:
                sys.stdout.write(
                    "{0}/{1} scripts written\r".format(tot_count, nscripts))
                sys.stdout.flush()

            fname = "{}.{}x.sh".format(os.path.join(path, script_obj.name), 1)
            with open(fname, 'w') as fp:
                fp.write(str(script_obj))

            tot_count += 1

        print "{0}/{0} scripts written".format(nscripts)

        return tot_count


class SlurmResourceManager(ResourceManager):
    def __init__(self, name, default_queue, unthrottled_limit=1000):
        super(SlurmResourceManager, self).__init__(name, default_queue, 
                                                   unthrottled_limit)

    def _submit(self, script_path):
        cmdline = 'sbatch -p {0}'.format(self.queue).split() + [script_path]
        rm_pipe = subprocess.Popen(cmdline, stdout=subprocess.PIPE, 
                                   stderr=subprocess.PIPE)
        
        stdout, stderr = rm_pipe.communicate()
                
        if rm_pipe.returncode != 0:
            msg = "Error submitting jobs through `{0}'".format(self.name)
            msg = "Resource manager output:"
            msg += "\n\tstdout: `{0}'".format(stdout)
            msg += "\n\tstderr: `{0}'".format(stderr)
            
            raise JobSubmitFail(msg)

    def running_job_count(self):
        username = getpass.getuser()

        cmdargs = 'squeue -a -u {0}'.format(username)

        squeue_pipe = subprocess.Popen(cmdargs.split(), stdout=subprocess.PIPE,
                                       stderr=subprocess.PIPE)

        stdout, stderr = squeue_pipe.communicate()

        if squeue_pipe.returncode != 0:
            msg = "Subprocess call failed -- `{0}': ".format(cmdargs)
            msg += str(stderr) + '\n'
            raise JobCountFailed(msg)

        lines = [line for line in stdout.split('\n') if username in line]

        return len(lines)
    
    def add_job(self, name, command, opts):
        self._joblist.append(jobscripts.SlurmScript(name, command, opts))




resource_managers = {'olympus': SlurmResourceManager('olympus', 'shared'),
                     }

resource_managers['default'] = resource_managers['olympus']


def set_default_rm(name):
    if name in resource_managers:
        resource_managers['default'] = resource_managers['sbatch']
    else:
        raise KeyError("Unrecognized resource manager name `{0}'".format(name))

 
